// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief Main Body of the PCMH MPI program
*/

#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "SMG_All.h"

struct option long_options[] = {
	{ "bmap", required_argument, NULL, 'm' }, // Path to the map file.
	{ "samples", required_argument, NULL, 's' }, // Path to the map file.
	{ "sped", required_argument, NULL, 'p' }, // Path to the PED file.
	{ "id", required_argument, NULL, 'i' }, // Run ID
	{ "rdir", required_argument, NULL, 'r'}, // Remote Output direct
	{ "ldir", required_argument, NULL, 'l'}, // Local Output direct
	{ "perm", required_argument, NULL, 'P'}, // Path to the permutation file.
	{ "merge", optional_argument, NULL, 'M'}, // Flag about to do the merge.
	{ "missed-rank", required_argument, NULL, 'n'}, // A rank that failed for whatever reason that needs re-running.
	{ "missed-procs", required_argument, NULL, 'o'}, // No. processors associated with the original calculation.
	{ "sol", required_argument, NULL, 'S' }, // Strata for layer file.
	{ "ci", required_argument, NULL, 'c' }, // Confidence interval
	{ "sort-by", required_argument, NULL, 't' }, // Column Name to sort final output by.
	{ 0,    0,    0,    0   } 
};

int pcmh_bt_main(int argc, char *argv[]) {
	char ch, *spedData = NULL;
	struct selected_options ops;
	int myrank, numprocs, result;
	BOOL jobDone = FALSE;
	struct chromosomes chromos;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];
	int nArgCount = 0, nRequiredArgs = 8, i, j;
	struct small_locus *map = NULL;
	struct sample *samples = NULL;
	int nMap = 0, nSamples = 0, nSpedSamples = 0, *pPermutations = NULL;
	int nSNPs = 0, nSpedRecordSizeInBytes = 0;
	BOOL isDoingMerge = FALSE;
	int perm_rows = 0, perm_columns = 0;
	double *originals = NULL, *results = NULL, zt;
	int *R = NULL, *maxR = NULL, missedRank = -1, missedProcs = -1, nAdj = 0;
	struct adjusted *adj = NULL;

	// Initialise MPI environment, get size and rank.
	if ((result = MPI_Init(&argc,&argv)) != MPI_SUCCESS) mpi_error(result, "Initialisation");
	if ((result = MPI_Comm_size(MPI_COMM_WORLD,&numprocs)) != MPI_SUCCESS) mpi_error(result, "Number of processors");
	if ((result = MPI_Comm_rank(MPI_COMM_WORLD,&myrank)) != MPI_SUCCESS) mpi_error(result, "Processor number");

	init_genrand((unsigned long) time(NULL));
	options_init(&ops);
	chromosomes_init(&chromos);
	chromosomes_define_human(&chromos);

	// Want to set these set via the command-line.
	*ops.ldir = *ops.rdir = *ops.szRunId = '\0'; 
	ops.verbose = FALSE;

	while ((ch = (char) getopt_long(argc, argv, "m:v:p:f:i:", long_options, NULL)) != -1) {
		double value = 0.0;
		switch (ch) {
			case 'i':
				strcpy(ops.szRunId, optarg);
				nArgCount++;
				break;
			case 's': // Sample File
				strcpy(ops.szSampleFileName, optarg);
				nArgCount++;
				break;
			case 'p': // SPED file
				strcpy(ops.szSpedFileName, optarg);
				nArgCount++;
				break;
			case 'm': // Bmap File
				strcpy(ops.szBmapFileName, optarg);
				nArgCount++;
				break;
			case 'r': // Remote output directory
				strcpy(ops.rdir, optarg);
				nArgCount++;
				break;
			case 'l': // Local output directory
				strcpy(ops.ldir, optarg);
				nArgCount++;
				break;
			case 'P': // Path to permutation file.
				strcpy(ops.szPermFileName, optarg);
				nArgCount++;
				break;
			case 'M':
				isDoingMerge = TRUE;
				break;
			case 'n': // Missed rank to re-run calculation on a single processor.
				missedRank = atoi(optarg);
				break;
			case 'o': // No. processors associated with a 'missed' rank.
				missedProcs = atoi(optarg);
				break;
			case 'S': // The strata/level file.
				strcpy(ops.szSolFileName, optarg);
				nArgCount++;
				break;
			case 'c': // The confidence interval
				value = atof(optarg);
				if (value >= 0.5 && value <= 1.0) {
					ops.ci_level = value;
				}
				break;
			case 't': // Sort column name
				if (strlen(optarg) > 0) strcpy(ops.szSortColumnName, optarg);
				break;
			default:
				break;
		}
	}
	zt = ltqnorm(1 - (1 - ops.ci_level) / 2); // Z-value of the CI.
	
	// If the required arguements are not specified, there's no point carrying on.
	if (nArgCount != nRequiredArgs) {
		chromosomes_free_attributes(&chromos);
		MPI_Finalize();
		return -1;
	}
	
	// See if a calculation needs to be re-run on a 'single' processor.
	// Check that all the command-line options are set.
	if (missedRank >= 0 || missedProcs >= 0) {
		BOOL exitEarly = FALSE;

		if (numprocs != 1) exitEarly = TRUE;
		else if (missedRank == -1 || missedProcs == -1) exitEarly = TRUE;
		else if (!(missedRank >= 0 && missedRank <= missedProcs)) exitEarly = TRUE;

		if (exitEarly) {
			MPI_Finalize();
			return -1;
		} else {
			myrank = missedRank;
			numprocs = missedProcs;
		}
	}
	
	// Do the permuted calculations.
	if (!isDoingMerge) {
		if ((map = bmapfile_load(&ops, &nMap)) != NULL && 
			(samples = samplefile_load(&ops, &nSamples)) != NULL &&
			(pPermutations = permfile_load(&ops, &perm_rows, &perm_columns)) != NULL &&
			(spedData = spedfile_load(&ops, &nSNPs, &nSpedSamples, &nSpedRecordSizeInBytes)) != NULL &&
			solfile_load(&ops, samples, nSamples)) 
		{
			if (nMap > 0 && perm_rows > 0 && perm_columns == nSamples && nSNPs == nMap && ops.nk > 1) {
				BOOL adjAllocated = TRUE;
				int nPermsPerProcessor = get_tasks_per_processor(numprocs, perm_rows);
				int start_perm, end_perm;

				start_perm = myrank * nPermsPerProcessor;
				end_perm = start_perm + nPermsPerProcessor;
				if (end_perm > perm_rows || myrank == numprocs - 1) end_perm = perm_rows;

				originals = create_vector(nSNPs, 0);
				results = create_vector(nSNPs, 0);
				R = create_int_array(nSNPs, 0);
				maxR = create_int_array(nSNPs, 0);
				if (myrank == ROOT) {
					adj = adjusted_create_array(nSNPs);
					if (adj == NULL) adjAllocated = FALSE;
				}

				// Permute the data-set.
				if (originals != NULL && R != NULL && maxR != NULL && results != NULL && adjAllocated) {
					samplefile_aff_coding(samples, nSamples);
					samplefile_set_original_order(samples, nSamples);
					samplefile_set_missing_by_sol(samples, nSamples);
					
					// Non-permuted calculation.
					for (i = 0; i < nSNPs; i++) {
						BOOL X = FALSE, haploid = FALSE;
						BitString bs = (BitString) (spedData + (i * nSpedRecordSizeInBytes));

						// Check if the SNP is located on autosome.
						if (chromos.chr_sex[map[i].chr]) X = TRUE;
						else if (chromos.chr_haploid[map[i].chr]) haploid = TRUE;

						// Calculate the original pvalue and affected counts.
						if (myrank != ROOT) originals[i] = pcmh_bt_assoc(zt, ops.nk, X, haploid, bs, samples, nSamples, NULL, NULL, 0, 0, ops.breslowday, NULL);
						else originals[i] = pcmh_bt_assoc(zt, ops.nk, X, haploid, bs, samples, nSamples, NULL, NULL, 0, 0, ops.breslowday, &adj[i]);
					}

					// Do the permutations for this rank.
					for (i = start_perm; i < end_perm && i < perm_rows; i++) {
						long offset = (long) i * (long) perm_columns;
						int *pPerms = pPermutations + offset;
						
						zero_vector(results, nSNPs);
						samplefile_permute_phenotype(samples, pPerms, nSamples);
						for (j = 0; j < nSNPs; j++) {
							BOOL X = FALSE, haploid = FALSE;
							BitString bs = (BitString) (spedData + (j * nSpedRecordSizeInBytes));

							// Check if the SNP is located on autosome.
							if (chromos.chr_sex[map[j].chr]) X = TRUE;
							else if (chromos.chr_haploid[map[j].chr]) haploid = TRUE;

							// Calculate the original pvalue and affected counts.
							results[j] = pcmh_bt_assoc(zt, ops.nk, X, haploid, bs, samples, nSamples, NULL, NULL, 0, 0, ops.breslowday, NULL);
						}
						perm_update(nSNPs, results, originals, R, maxR);
					}

					jobDone = count_file_write_and_copy(&ops, SUFFIX_CMH_FILE, myrank, numprocs, abs(end_perm - start_perm), nSNPs, R, maxR) ? TRUE : FALSE;
					if (myrank == ROOT && jobDone) jobDone = adjusted_write_and_copy(adj, nSNPs, &ops, myrank) ? TRUE : FALSE;
				}
			}
		}
	} else {
		// Merging & sorting the output files.
		if (numprocs == 1) { // This code in this mode needs to run as a singleton.
			BOOL *file_copy_failures = NULL, hasCopiedAll = FALSE;
			char line[MAX_LINE_LENGTH];

			map = bmapfile_load(&ops, &nMap);
			spedData = spedfile_load(&ops, &nSNPs, &nSpedSamples, &nSpedRecordSizeInBytes);
			samples = samplefile_load(&ops, &nSamples);
			numprocs = readNumberOfProcs(remoteFilepath, line, MAX_LINE_LENGTH, ops.rdir, ops.szRunId, SUFFIX_CMH_FILE);
			result = solfile_load(&ops, samples, nSamples);
			adj = _adjusted_load(&ops, &nAdj);

			if (map != NULL && nMap > 0 && spedData != NULL && nSNPs > 0 && nSNPs == nMap && samples != NULL && nSamples > 0 && numprocs > 0 && result &&
				adj != NULL && nAdj > 0 && nAdj == nSNPs) 
			{
				int nPermutations = 0;
				file_copy_failures = create_bool_array(numprocs, FALSE); 
				R = create_int_array(nSNPs, 0);
				maxR = create_int_array(nSNPs, 0);

				if(count_file_copy_and_parse(&ops, SUFFIX_CMH_FILE, numprocs, nSNPs, &nPermutations, R, maxR, &hasCopiedAll, file_copy_failures)) {
					// Create the local output file
					if (hasCopiedAll) {
						FILE *out = NULL;

						sprintf(localFilepath, "%s%s%s", ops.ldir, ops.szRunId, SUFFIX_CMH_FILE);
						sprintf(remoteFilepath, "%s%s%s", ops.rdir, ops.szRunId, SUFFIX_CMH_FILE);
						out = fopen(localFilepath, "w");
						if (out != NULL) {
							struct adjusted *a = NULL; 
							double nperms = (double) nPermutations + 1;

							samplefile_aff_coding(samples, nSamples);
							samplefile_set_original_order(samples, nSamples);
							samplefile_set_missing_by_sol(samples, nSamples);

							// Write out the results of each SNP to file.
							pcmh_bt_assoc_write_file_header(out, &ops);
							for (i = 0, a = adj; i < nSNPs; i++, a++) {
								double emp1, emp2;
								BOOL X = FALSE, haploid = FALSE;
								BitString bs = (BitString) (spedData + (i * nSpedRecordSizeInBytes));

								// Check if the SNP is located on autosome.
								if (chromos.chr_sex[map[i].chr]) X = TRUE;
								else if (chromos.chr_haploid[map[i].chr]) haploid = TRUE;

								// Ratio counts
								emp1 = (double) (R[i] + 1) / nperms;
								emp2 = (double) (maxR[i] + 1) / nperms;

								pcmh_bt_assoc(zt, ops.nk, X, haploid, bs, samples, nSamples, &map[i], out, emp1, emp2, ops.breslowday, a);
							}
							fclose(out);

							jobDone = sort_file_by_col_and_copy(localFilepath, remoteFilepath, ops.szSortColumnName) ? TRUE : FALSE;
							if (jobDone) assoc_cleanup_temp_files(&ops, SUFFIX_CMH_FILE, numprocs);
						}
					}
				} else {
					write_failed_rank_log(&ops, numprocs, file_copy_failures);
				}
			}
			if (file_copy_failures != NULL) free(file_copy_failures);
		} 
	}

	// Free now as not needed anymore
	if (pPermutations != NULL) free(pPermutations);
	if (samples != NULL) free(samples);
	if (spedData != NULL) free(spedData);
	if (map != NULL) free(map); 
	if (originals != NULL) free(originals);
	if (R != NULL) free(R);
	if (maxR != NULL) free(maxR);
	if (results != NULL) free(results);
	if (adj != NULL) free(adj);
	chromosomes_free_attributes(&chromos); 

	MPI_Finalize();

	if (jobDone) return 0;
	else return -1;
}

int main(int argc, char *argv[]) {
	return pcmh_bt_main(argc, argv);
}